import numpy as np
import matplotlib
matplotlib.use('TkAgg')  # 或 'Qt5Agg'
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.font_manager as fm

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

# 尝试其他可能的中文字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS', 'STHeiti']


def calculate_optimized_camera_position(serve_pos, landing_pos, initial_distance=12.0, height_offset=5.0,
                                       angle_offset=30, trajectory_height=3.0):
    """
    计算优化后的相机初始位置，使其与球轨迹有一个合适的夹角，并能一镜到底。

    参数:
    - serve_pos: 发球点坐标 (x, y, z)
    - landing_pos: 落点坐标 (x, y, z)
    - initial_distance: 相机与落点的初始水平距离
    - height_offset: 相机的高度偏移
    - angle_offset: 相机与球轨迹的夹角 (度)
    - trajectory_height: 球轨迹的最大高度，用于动态调整相机的高度

    返回:
    - 相机初始位置和观察点
    """
    serve = np.array(serve_pos, dtype=float)
    landing = np.array(landing_pos, dtype=float)

    # 计算从发球点到落点的方向向量
    direction = landing - serve
    direction_xy = np.array([direction[0], direction[1], 0])
    norm_xy = np.linalg.norm(direction_xy)

    if norm_xy == 0:
        direction_xy_unit = np.array([1, 0, 0])
    else:
        direction_xy_unit = direction_xy / norm_xy

    # 计算与球轨迹形成夹角的旋转矩阵
    angle_rad = np.radians(angle_offset)
    rotation_matrix = np.array([
        [np.cos(angle_rad), -np.sin(angle_rad)],
        [np.sin(angle_rad), np.cos(angle_rad)]
    ])

    # 旋转方向向量，得到新的相机方向
    rotated_direction = np.zeros(3)
    rotated_direction[:2] = rotation_matrix.dot(direction_xy_unit[:2])

    # 假设球的轨迹是一个抛物线，基于落点与发球点的高度差和最大轨迹高度，推算相机的高度
    height_difference = landing[2] - serve[2]
    camera_height_adjustment = height_offset + trajectory_height * 1.2  # 确保相机高度始终高于球轨迹

    # 计算相机的位置，确保相机的初始位置足够远，且保持一定的高度偏移
    camera_pos = landing + rotated_direction * initial_distance + np.array([0, 0, camera_height_adjustment])

    # 观察点设为球的轨迹中点，给相机一个稳定的观察目标
    look_at = (serve + landing) / 2

    return tuple(camera_pos), tuple(look_at)


def generate_camera_trajectory(serve_pos, landing_pos, ball_trajectory, max_height=3.0):
    """
    根据球的轨迹生成相机的移动轨迹
    
    参数:
    - serve_pos: 发球点坐标
    - landing_pos: 落点坐标
    - ball_trajectory: 球的轨迹点数组
    - max_height: 球轨迹的最大高度
    
    返回:
    - 相机位置和观察点的数组
    """
    serve = np.array(serve_pos, dtype=float)
    landing = np.array(landing_pos, dtype=float)
    
    # 计算初始相机位置，传入轨迹最大高度
    initial_pos, _ = calculate_optimized_camera_position(
        serve, landing, 
        initial_distance=15.0, 
        angle_offset=30,
        trajectory_height=max_height
    )
    
    # 最终相机位置：落点上方，高度基于球轨迹高度
    final_pos = np.array(landing) + np.array([0, 0, max(6.0, max_height * 1.5)])
    
    initial_pos = np.array(initial_pos)
    
    camera_positions = []
    look_at_points = []
    
    for i, ball_pos in enumerate(ball_trajectory):
        # 计算进度
        progress = i / (len(ball_trajectory) - 1)
        
        # 使用easeInOutQuad缓动函数使移动更自然
        if progress < 0.5:
            t = 2 * progress * progress
        else:
            t = -1 + (4 - 2 * progress) * progress
        
        # 相机位置从初始位置逐渐移动到最终位置
        camera_pos = initial_pos * (1 - t) + final_pos * t
        
        # 相机的观察点始终跟随球的位置
        # 早期更关注球，后期更关注落点
        ball_weight = max(0.0, 1.0 - progress * 1.5)
        look_at = ball_pos * ball_weight + np.array(landing) * (1 - ball_weight)
        
        camera_positions.append(tuple(camera_pos))
        look_at_points.append(tuple(look_at))
    
    return camera_positions, look_at_points


def generate_trajectory(serve_pos, landing_pos, num_points=100, max_height=3.0):
    p1 = np.array(serve_pos)
    p2 = np.array(landing_pos)
    t = np.linspace(0, 1, num_points)
    # 抛物线模拟
    trajectory = (1 - t)[:, None] * p1 + t[:, None] * p2
    trajectory[:, 2] += max_height * np.sin(np.pi * t)  # 高度峰值在中间
    return trajectory


def draw_tennis_court(ax):
    """绘制网球场地"""
    # 场地边界
    court_length = 23.77  # 全场长度
    court_width = 8.23    # 单打场地宽度
    
    # 中心线
    ax.plot([court_length/2, court_length/2], [-court_width/2, court_width/2], [0, 0], 'k-', alpha=0.7)
    
    # 绘制球网 (高度1.07米)
    net_height = 1.07
    net_x = np.array([court_length/2, court_length/2])
    net_y = np.array([-court_width/2-1, court_width/2+1])
    net_z_bottom = np.array([0, 0])
    net_z_top = np.array([net_height, net_height])
    ax.plot(net_x, net_y, net_z_bottom, 'k-', linewidth=2)
    ax.plot(net_x, net_y, net_z_top, 'k-', linewidth=2)
    
    # 绘制球场边界
    court_corners_x = [0, court_length, court_length, 0, 0]
    court_corners_y = [-court_width/2, -court_width/2, court_width/2, court_width/2, -court_width/2]
    court_corners_z = [0, 0, 0, 0, 0]
    ax.plot(court_corners_x, court_corners_y, court_corners_z, 'k-', alpha=0.7)
    
    # 发球区
    service_line = 6.4  # 发球线距离底线的距离
    ax.plot([service_line, service_line], 
            [-court_width/2, court_width/2], [0, 0], 'k--', alpha=0.5)
    ax.plot([court_length-service_line, court_length-service_line], 
            [-court_width/2, court_width/2], [0, 0], 'k--', alpha=0.5)
    
    # 中央发球线
    ax.plot([service_line, court_length-service_line], [0, 0], [0, 0], 'k--', alpha=0.5)


def test_camera_position():
    """测试不同场景下的相机位置是否合理"""
    # 定义多组测试数据 (发球点和落点)
    test_scenarios = [
        {
            "name": "右场发球到左场底线",
            "serve": (3, 4, 1.5),
            "landing": (20, -3, 0),
            "max_height": 2.5
        },
        {
            "name": "右场底线抽球到左场中场",
            "serve": (1, 3, 1.0),
            "landing": (15, -2, 0),
            "max_height": 1.8
        },
        {
            "name": "左场发球到右场底线角落",
            "serve": (3, -4, 1.5),
            "landing": (20, 3.5, 0),
            "max_height": 2.5
        },
        {
            "name": "左场底线到右场发球区",
            "serve": (1, -3, 1.0),
            "landing": (18, 2, 0),
            "max_height": 1.5
        },
        {
            "name": "网前截击",
            "serve": (11, 3, 1.2),
            "landing": (15, -2, 0),
            "max_height": 0.8
        }
    ]
    
    # 为每个场景计算并显示相机位置
    fig = plt.figure(figsize=(15, 12))
    
    for i, scenario in enumerate(test_scenarios):
        ax = fig.add_subplot(3, 2, i+1, projection='3d')
        
        serve = scenario["serve"]
        landing = scenario["landing"]
        max_height = scenario["max_height"]
        
        # 计算轨迹
        ball_trajectory = generate_trajectory(serve, landing, max_height=max_height)
        
        # 计算相机位置，传入轨迹最大高度
        initial_camera, look_at = calculate_optimized_camera_position(
            serve, landing, 
            initial_distance=15.0, 
            height_offset=5.0, 
            angle_offset=30,
            trajectory_height=max_height
        )
        
        # 计算相机轨迹，传入轨迹最大高度
        camera_positions, _ = generate_camera_trajectory(serve, landing, ball_trajectory, max_height)
        final_camera = camera_positions[-1]
        
        # 绘制球场
        draw_tennis_court(ax)
        
        # 绘制球的轨迹
        ax.plot(ball_trajectory[:, 0], ball_trajectory[:, 1], ball_trajectory[:, 2], 'r-', alpha=0.6)
        
        # 绘制发球点和落点
        ax.scatter(*serve, color='green', s=50, label='发球点')
        ax.scatter(*landing, color='blue', s=50, label='落点')
        
        # 绘制相机初始位置和最终位置
        ax.scatter(*initial_camera, color='orange', s=80, marker='^', label='相机初始位置')
        ax.scatter(*final_camera, color='purple', s=80, marker='*', label='相机最终位置')
        
        # 添加相机坐标文本
        ax.text(initial_camera[0], initial_camera[1], initial_camera[2] + 0.5, 
                f'初始: ({initial_camera[0]:.1f}, {initial_camera[1]:.1f}, {initial_camera[2]:.1f})', 
                color='orange', fontsize=8)
        ax.text(final_camera[0], final_camera[1], final_camera[2] + 0.5, 
                f'最终: ({final_camera[0]:.1f}, {final_camera[1]:.1f}, {final_camera[2]:.1f})', 
                color='purple', fontsize=8)
        
        # 绘制相机观察线
        ax.plot([initial_camera[0], look_at[0]], 
                [initial_camera[1], look_at[1]], 
                [initial_camera[2], look_at[2]], 'k--', alpha=0.4)
        
        # 设置视图
        ax.view_init(elev=30, azim=-60)
        ax.set_xlim(0, 24)
        ax.set_ylim(-10, 10)
        ax.set_zlim(0, 8)
        
        ax.set_title(f"场景: {scenario['name']}")
        ax.set_xlabel('X轴')
        ax.set_ylabel('Y轴')
        ax.set_zlabel('Z轴')
        
        if i == 0:  # 只在第一个子图显示图例
            ax.legend(loc='upper left')
    
    plt.tight_layout()
    plt.savefig('camera_position_test.png', dpi=300)
    plt.show()
    
    print("相机位置测试结果分析:")
    for i, scenario in enumerate(test_scenarios):
        serve = scenario["serve"]
        landing = scenario["landing"]
        max_height = scenario["max_height"]
        
        # 传入轨迹最大高度
        initial_camera, _ = calculate_optimized_camera_position(
            serve, landing, 
            initial_distance=15.0, 
            height_offset=5.0, 
            angle_offset=30,
            trajectory_height=max_height
        )
        
        # 计算相机到发球点和落点的距离
        dist_to_serve = np.linalg.norm(np.array(initial_camera) - np.array(serve))
        dist_to_landing = np.linalg.norm(np.array(initial_camera) - np.array(landing))
        
        # 计算相机高度与最大球高度的比例
        max_ball_height = max_height + max(serve[2], landing[2])
        height_ratio = initial_camera[2] / max_ball_height
        
        print(f"\n{i+1}. 场景: {scenario['name']}")
        print(f"   发球点: {serve}, 落点: {landing}")
        print(f"   球轨迹最大高度: {max_height:.2f}米")
        print(f"   相机初始位置: {initial_camera}")
        print(f"   相机到发球点距离: {dist_to_serve:.2f}米")
        print(f"   相机到落点距离: {dist_to_landing:.2f}米")
        print(f"   相机高度与最大球高度比例: {height_ratio:.2f}")
        print(f"   评估: {'合理' if (dist_to_landing > 8 and height_ratio > 1.2) else '需要调整'}")


def main():
    # 测试相机位置
    test_camera_position()
    
    # 以下是原始的单一场景动画展示
    # 参数设置：发球点和落点
    serve = (3, 3, 1.5)
    landing = (20, -2, 0)
    max_height = 2.5
    
    # 生成球的轨迹
    ball_trajectory = generate_trajectory(serve, landing, max_height=max_height)
    
    # 生成相机轨迹，传入轨迹最大高度
    camera_positions, look_at_points = generate_camera_trajectory(serve, landing, ball_trajectory, max_height)
    
    # 设置图形
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 设置相机视角（初始）
    ax.view_init(elev=30, azim=-60)
    ax.set_xlim(0, 24)
    ax.set_ylim(-8, 8)
    ax.set_zlim(0, 8)
    
    # 标签和网球场地参考
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('业余比赛网球鹰眼挑战系统')
    
    # 绘制球场
    draw_tennis_court(ax)
    
    # 静态元素：发球点、落点
    ax.scatter(*serve, color='red', s=50, label='发球点')
    ax.scatter(*landing, color='blue', s=50, label='落点')
    ax.legend()
    
    # 动态球对象
    ball, = ax.plot([], [], [], 'o', color='orange', markersize=8)
    
    # 相机位置指示器
    camera, = ax.plot([], [], [], 'o', color='green', markersize=6, label='相机位置')
    camera_line, = ax.plot([], [], [], 'g--', alpha=0.5)  # 从相机到观察点的线
    
    # 添加球的轨迹线（淡化显示）
    ax.plot(ball_trajectory[:, 0], ball_trajectory[:, 1], ball_trajectory[:, 2], 
           'r--', alpha=0.2)
           
    # 添加相机坐标文本（初始为空，会在动画中更新）
    camera_text = ax.text2D(0.02, 0.95, '', transform=ax.transAxes, fontsize=10)
    
    # 动画函数
    def animate(i):
        # 更新球位置
        ball.set_data([ball_trajectory[i, 0]], [ball_trajectory[i, 1]])
        ball.set_3d_properties([ball_trajectory[i, 2]])
        
        # 更新相机位置
        camera_pos = camera_positions[i]
        look_at = look_at_points[i]
        
        camera.set_data([camera_pos[0]], [camera_pos[1]])
        camera.set_3d_properties([camera_pos[2]])
        
        # 更新相机到观察点的线
        camera_line.set_data([camera_pos[0], look_at[0]], [camera_pos[1], look_at[1]])
        camera_line.set_3d_properties([camera_pos[2], look_at[2]])
        
        # 更新相机坐标文本
        camera_text.set_text(f'相机坐标: ({camera_pos[0]:.2f}, {camera_pos[1]:.2f}, {camera_pos[2]:.2f})')
        
        return ball, camera, camera_line, camera_text
    
    # 创建动画
    ani = FuncAnimation(fig, animate, frames=len(ball_trajectory), interval=50, blit=True)
    
    plt.show()


if __name__ == "__main__":
    main()
